﻿module BernulliMixture
open MathNet.Numerics
open MathNet.Numerics.FSharp
open MathNet.Numerics.LinearAlgebra.Double
open System
open System.Collections.Generic
open Types
open Utils


//compute xn given uk
let p_xn_given_uk (xn : DenseVector) (uk : DenseVector) =
  let mutable acc = 1.

  for i in [0..xn.Count-1] do
    let first = 
      if xn.[i] = 0. then 1.
      else  uk.[i]

    let second = 
      if xn.[i] = 1. then 1.
      else  1. - uk.[i]
    acc <- acc * first * second
  acc

//E optimization
let compute_probabilities (xn_list : DenseVectorWrapper list) (mu_k_list : DenseVectorWrapper list) =
  for xn in xn_list do
    for mu in mu_k_list do
      let p_xn_given_uk = p_xn_given_uk xn.DenseVector mu_k_list.[mu.Index].DenseVector
      xn.PK.[mu.Index] <- p_xn_given_uk

let gamma_znk (pi : float list) (mu : DenseVectorWrapper list) (k : int) (xn : DenseVectorWrapper) =
  let first = pi.[k] * xn.PK.[k]
  if first = 0. then 0.
  else 
    let mutable acc = 0.
    for i in [0..mu.Length-1] do
      acc <- acc + pi.[i] * xn.PK.[mu.[i].Index]
    first / acc

//M optimization
let n_k (k_gamma_list : DenseMatrix) (k : int) (xn : DenseVectorWrapper list) =
  let mutable acc = 0.
  for x in xn do 
    acc <- acc + k_gamma_list.[x.Index, k]
  acc  

let pi_k (k_gamma : DenseMatrix) (k : int) (n : int) (xn : DenseVectorWrapper list) =
  n_k k_gamma k xn / (float n)

let u_k (k_gamma_list : DenseMatrix) (mu : DenseVectorWrapper list) (k : int) (xn : DenseVectorWrapper list) : DenseVectorWrapper =
  let mutable acc = DenseVector(xn.Head.DenseVector.Count)
  let res = 
    for x in xn do 
      acc <- acc + (k_gamma_list.[x.Index, k] * x.DenseVector)
    acc  
  let res = res / n_k k_gamma_list k xn
  let res = DenseVectorWrapper(res, k, k, k |> ref, mu.[k].PK)
  res

//compute R
let compute_r (input : DenseVectorWrapper list) =
  let mutable a = 0.
  let mutable b = 0.
  for p1 in input do
    for p2 in input do
      if p1.Index <> p2.Index then 
        if p1.PK.MaximumIndex() = p2.PK.MaximumIndex() && p1.RealNumber = p2.RealNumber then
          a <- a + 1.
        elif p1.PK.MaximumIndex() <> p2.PK.MaximumIndex() && p1.RealNumber <> p2.RealNumber then
          b <- b + 1.
  let r = (2. * (a + b)) / (float input.Length * (float input.Length - 1.))
  Console.WriteLine("R = " + string r)


